import os
import torch
from torch.utils.data import Dataset, DataLoader
import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt  # 新增导入

class Medical3DDataset(Dataset):
    def __init__(self, data_dir, num_classes=3):
        """
        初始化数据集类，支持多类别分割

        参数:
        data_dir (str): 数据根目录，包含 'data' 和 'label' 子目录
        num_classes (int): 类别数量，默认为3 (0,1,2)
        """
        self.data_dir = data_dir
        self.num_classes = num_classes
        self.data_path = os.path.join(data_dir, 'data')
        self.label_path = os.path.join(data_dir, 'label')
        self.data_files = sorted(os.listdir(self.data_path))
        self.label_files = sorted(os.listdir(self.label_path))

    def __len__(self):
        """
        返回数据集的长度。

        返回:
        int: 数据集中样本的数量。
        """
        return len(self.data_files)

    def __getitem__(self, idx):
        """
        根据索引获取数据样本和对应的标签。

        参数:
        idx (int): 样本的索引。

        返回:
        tuple: 包含数据样本和对应标签的元组。
        """
        data_file = os.path.join(self.data_path, self.data_files[idx])
        label_file = os.path.join(self.label_path, self.label_files[idx])

        # 加载原始数据
        data_nii = nib.load(data_file)
        data_array = data_nii.get_fdata()
        data_array = (data_array - np.mean(data_array)) / np.std(data_array)
        data_tensor = torch.tensor(data_array, dtype=torch.float32).unsqueeze(0)

        # 加载标签数据并转换为one-hot编码
        label_nii = nib.load(label_file)
        label_array = label_nii.get_fdata()

        # 创建one-hot编码标签
        label_tensor = torch.zeros((self.num_classes, *label_array.shape), dtype=torch.float32)
        for class_idx in range(self.num_classes):
            label_tensor[class_idx] = torch.tensor((label_array == class_idx).astype(np.float32))

        return data_tensor, label_tensor

    def show_slice(self, idx, slice_idx=None, class_idx=0):
        """
        显示指定数据的中间切片

        参数:
        idx (int): 数据索引
        slice_idx (int, optional): 指定切片索引，默认显示中间切片
        class_idx (int): 要显示的类别索引(0-背景,1-类别1,2-类别2)
        """
        data, label = self[idx]
        data = data.squeeze().numpy()  # 数据形状 [D, H, W]
        label = label[class_idx].numpy()  # 选择特定类别的标签 [D, H, W]

        if slice_idx is None:
            slice_idx = data.shape[0] // 2  # 默认显示中间切片

        plt.figure(figsize=(12, 6))

        plt.subplot(1, 2, 1)
        plt.imshow(data[slice_idx], cmap='gray')
        plt.title(f'Data Slice {slice_idx}')
        plt.axis('off')

        plt.subplot(1, 2, 2)
        plt.imshow(label[slice_idx], cmap='jet')
        plt.title(f'Class {class_idx} Label Slice {slice_idx}')
        plt.axis('off')

        plt.tight_layout()
        plt.show()

if __name__ == "__main__":
    data_dir = r'D:\dataset\traindata'
    dataset = Medical3DDataset(data_dir, num_classes=3)

    # 测试数据加载
    data, label = dataset[0]
    print(f"Data shape: {data.shape}")  # 应为 [1, D, H, W]
    print(f"Label shape: {label.shape}")  # 应为 [3, D, H, W] (3个类别)

    # 检查每个类别的像素数量
    for class_idx in range(3):
        print(f"Class {class_idx} pixels: {(label[class_idx] == 1).sum().item()}")

    # 显示第一个样本的中间切片
    dataset.show_slice(0)  # 这行应该放在这里而不是在show_slice方法内部

    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

    # 示例：遍历数据加载器
    for data, label in dataloader:
        print(f"Data shape: {data.shape}, Label shape: {label.shape}")

    # Check unique values in target labels
    unique_labels = set()
    for _, labels in dataloader:
        unique_labels.update(labels.unique().tolist())
    print(f"Unique labels in dataset: {unique_labels}")

